// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <melon/dynamic.h>

#include <pollux/common/base/simd_util.h>
#include <pollux/common/base/pollux_exception.h>
#include <pollux/vector/builder_type_utils.h>
#include <pollux/vector/simple_vector.h>

namespace kumo::pollux {
    /**
     * Provides a vector implementation for number types where the difference
     * between the maximum value and the minimum value is in a range that allows for
     * a smaller delta value to be encoded along with a base, bias value.  The
     * result is that we can store what might be logically a 64bit number in a space
     * that is a base 64bit value (the bias) and then smaller delta values, which
     * might be 32, 16 or 8 bit values.
     *
     * For example, we can encode the sequence of 64bit numbers:
     *   -1023, -963, -774
     * as 8bit numbers (since -774 - (-1023) = 249 which is < 255):
     *   bias = -898, values = -125, -65, 124
     *
     * or another example:
     *   1024, 1124, 900
     * as:
     *   bias = 1012, values = 12, 112, -112
     *
     * The original values at an index are retrieved via bias + value[i], e.g.
     * logical_value[1] = -898 + (-65) = -963
     *
     * To calculate the bias value we first find the delta value, which is the
     * difference between the max and min value in the vector. Since the final
     * biased values are stored in signed integer space, the bias needs to be
     * in between the max and min value of the vector. The biased values are
     * calculated by subtracting the bias value from input values. With the bias
     * in between min and max, the final values (input value - bias) will be
     * both positive and negative, utilizing the entire signed integer range.
     *
     * The bias is calculating using the formula -
     *  Delta  = maxValue - minValue
     *  Bias   = minValue + ceil(delta/2)
     *  Output = input - bias
     *
     * Example: [200, 250, 300, 420]
     *  Delta = 220 (420 - 200)
     *  Bias = 310 (200 + 220/2).
     *  Output = [-110, -60, -10, 110]
     *
     * The output values can be encoded as int8_t since they fit in the
     * range [-128, 127].
     *
     * The ceiling operation in the calculation is very important. We need
     * the delta to be higher since the absolute value of the smallest negative
     * value is higher than the largest positive value for any integer.
     * For instance, int8_t has the range [-128, 127] where abs(-128) > 127.
     *
     * Example: [500, 550, 400, 655]
     *  Delta  = 255 (655 - 400)
     *  Bias   = 528 (400 + ceil (255/2))
     *  Output = [-28,  22,  -128, 127]

     * If the bias value had been 527 instead of 528, 400 would be encoded
     * as -127 and 655 would be 128. This value (128) would not fit within int8_t
     * range.
     */
    template<typename T>
    class BiasVector : public SimpleVector<T> {
        // TODO(T54638002): Add support for unsigned types in biased vector
        // FIXME(venkatra): Commented out since this assertion is always true due
        // to a typo. If we fix it by changing 'admitsBias<T>' to 'admitsBias<T>()',
        // compilation fails due to lack of support for unsigned types.
        // static_assert(
        //     admitsBias<T>,
        //    "Only vectors that are stored as 64, 32 and 16-bit numbers can use
        //    bias");

    public:
        static constexpr bool can_simd =
        (std::is_same_v<T, int64_t> || std::is_same_v<T, int32_t> ||
         std::is_same_v<T, int16_t>);

        BiasVector(
            pollux::memory::MemoryPool *pool,
            BufferPtr nulls,
            size_t length,
            TypeKind valueType,
            BufferPtr values,
            T bias,
            const SimpleVectorStats<T> &stats = {},
            std::optional<int32_t> distinctCount = std::nullopt,
            std::optional<vector_size_t> nullCount = std::nullopt,
            std::optional<bool> sorted = std::nullopt,
            std::optional<ByteCount> represented_bytes = std::nullopt,
            std::optional<ByteCount> storageByteCount = std::nullopt);

        ~BiasVector() override {
        }

        bool contains_null_at(vector_size_t idx) const override {
            return BaseVector::is_null_at(idx);
        }

        const T valueAtFast(vector_size_t idx) const;

        const T value_at(vector_size_t idx) const override {
            SimpleVector<T>::checkElementSize();
            return valueAtFast(idx);
        }

        /**
         * Loads a SIMD vector of data at the virtual byteOffset given
         * Note this method is implemented on each vector type, but is intentionally
         * not virtual for performance reasons
         *
         * @param byteOffset - the byte offset to laod from
         */
        ksimd::batch<T> loadSIMDValueBufferAt(size_t index) const;

        std::unique_ptr<SimpleVector<uint64_t> > hash_all() const override;

        inline T bias() const {
            return bias_;
        }

        TypeKind valueType() const {
            return valueType_;
        }

        uint64_t retained_size() const override {
            return BaseVector::retained_size() + values_->capacity();
        }

        /**
         * Returns a shared_ptr to the underlying arrow array holding the values for
         * this vector. This is used during execution to process over the subset of
         * values when possible.
         */
        inline const BufferPtr &values() const override {
            return values_;
        }

        bool is_scalar() const override {
            return true;
        }

        VectorPtr slice(vector_size_t, vector_size_t) const override {
            POLLUX_NYI();
        }

        VectorPtr copy_preserve_encodings(
            pollux::memory::MemoryPool *pool = nullptr) const override {
            auto selfPool = pool ? pool : BaseVector::pool_;
            return std::make_shared<BiasVector<T> >(
                selfPool,
                AlignedBuffer::copy(selfPool, BaseVector::nulls_),
                BaseVector::length_,
                valueType_,
                AlignedBuffer::copy(selfPool, values_),
                bias_,
                SimpleVector<T>::stats_,
                BaseVector::distinctValueCount_,
                BaseVector::nullCount_,
                SimpleVector<T>::isSorted_,
                BaseVector::representedByteCount_,
                BaseVector::storageByteCount_);
        }

    private:
        template<typename U>
        inline ksimd::batch<T> loadSIMDInternal(size_t byteOffset) const {
            auto mem = reinterpret_cast<const U *>(
                rawValues_ + byteOffset / sizeof(T) * sizeof(U));
            return ksimd::batch<T>::load_unaligned(mem);
        }

        TypeKind valueType_;
        BufferPtr values_;
        const uint8_t *rawValues_;

        // Note: there is no 64 bit internal array as the largest number type we
        // support is 64 bit and all biasing requires a smaller internal type.
        T bias_;

        // Used to debias several values at a time.
        std::conditional_t<can_simd, ksimd::batch<T>, char> biasBuffer_;
    };

    template<typename T>
    using BiasVectorPtr = std::shared_ptr<BiasVector<T> >;
} // namespace kumo::pollux

#include <pollux/vector/bias_vector-inl.h>
